import numpy as np
import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore.ops import functional as F
from mindspore.common.initializer import initializer


class TestNet(nn.Cell):
    def __init__(self):
        super(TestNet, self).__init__()
        self.op = ops.NLLLoss()

    def construct(self, input_x, labels, weight):
        m = nn.LogSoftmax(axis=1)
        out, out_weight = self.op(m(input_x), labels, weight)
        return [out]


if __name__ == "__main__":
    net = TestNet()

    input_x = ms.Tensor(np.ones([3,5]), ms.float32)
    labels = ms.Tensor(np.random.random_sample([3]), ms.int32)
    weight = ms.Tensor(np.ones([5]), ms.float32)
    result = net(input_x, labels, weight)
    print(result[0])
    print(result[0].shape)
